#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2021/3/4 9:52
# @Author  : CHEN Wang
# @Site    : 
# @File    : DigQuant.py
# @Software: PyCharm 

"""
脚本说明: 获取点宽数据
"""
import datetime
import numpy
import pandas as pd
import requests
import json
from quant_researcher.quant.project_tool.pretty_table import QA_util_to_json_from_pandas
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.time_tool import QA_util_date_valid
from quant_researcher.quant.project_tool.QACode import QA_util_code_tolist
from quant_researcher.quant.datasource_fetch.DataStruct import QA_DataStruct_Stock_day, QA_DataStruct_Stock_min


def get_dqt_data_by_api(data_type, input_data):
    """
    根据数据类型和传入参数获取点宽的数据

    :param str data_type: 数据类型，格式形如quote/stock/kdata/day，可以获取哪些数据见
                          http://192.168.90.203:19509/swagger-ui.html
    :param dict input_data: 输入的参数
    :return:
    """
    url = f'http://192.168.90.203:19509/{data_type}'

    input_data = json.dumps(input_data)

    response = requests.post(url, data=input_data, headers={"Content-Type": "application/json; charset=UTF-8"})

    if response.status_code == 200:
        response = response.json()
        if response['errInfo']['errCode'] == 0:
            if len(response['data']) != 0:
                return response['data']
            else:
                return None
        else:
            return None
    else:
        return None


def QA_fetch_stock_day_adv(
        code,
        start='all',
        end=None,
        if_drop_index=True,
):
    """
    :param code:  股票代码列表
    :param start: 开始日期,如果start=='all',则获取从1990-01-01至今的数据
    :param end:   结束日期
    :param if_drop_index:
    :return: 如果股票代码不存 或者开始结束日期不存在 在返回 None ，合法返回 QA_DataStruct_Stock_day 数据
    """
    # 获取股票日线数据并返回QA_DataStruct_Stock_day类型数据
    end = start if end is None else end
    start = str(start)[0:10]
    end = str(end)[0:10]

    if start == 'all':
        start = '1990-01-01'
        end = str(datetime.date.today())

    res = QA_fetch_stock_day(code, start, end, format='pd')
    if res is None:
        print(
            "DigQuant Error QA_fetch_stock_day_adv parameter code=%s , start=%s, end=%s call QA_fetch_stock_day return None"
            % (code,
               start,
               end)
        )
        return None
    else:
        res_reset_index = res.set_index(['date', 'code'], drop=if_drop_index)
        return QA_DataStruct_Stock_day(res_reset_index)


def QA_fetch_stock_day(
        code,
        start,
        end,
        format='numpy'
):
    """'获取股票日线数据并返回，默认返回numpy格式，支持输出np,pd,dict,list,json格式'
    Returns:
        [type] -- [description]

    """

    start = str(start)[0:10]
    end = str(end)[0:10]
    code = QA_util_code_tolist(code)  # 代码补齐6位数，并转换成list
    mkt_code = ['sse' if x[0] == '6' else 'szse' for x in code]
    code_df = pd.DataFrame({'code': code, 'market': mkt_code})
    code_dict = code_df.to_dict(orient='records')

    if QA_util_date_valid(end):  # 判断end的日期字符串格式是否为%Y-%m-%d
        res = get_dqt_data_by_api(data_type='quote/stock/kdata/day',
                                  input_data={
                                      "dateDurationTO": {
                                          "beginDate": start,
                                          "endDate": end
                                      },
                                      "filledUp": True,
                                      "marketCodeTOS": code_dict,
                                      "restorationWay": "NA"
                                  })

        if res is None:
            pass
        else:
            res_list = []
            for tmp in res:
                res_list.extend(tmp['stockKDays'])
            res = pd.DataFrame(res_list)
            res = res[['code', 'tradeDate', 'open', 'high', 'low', 'close', 'volume', 'turnOver']]
            res = res.rename(columns={'tradeDate': 'date', 'turnOver': 'amount'})
            res.date = pd.to_datetime(res.date)
            res = res.set_index('date', drop=False)
            res = res.loc[:, ['code', 'open', 'high', 'low', 'close', 'volume', 'amount', 'date']]

        if format in ['P', 'p', 'pandas', 'pd']:
            return res
        elif format in ['json', 'dict']:
            return QA_util_to_json_from_pandas(res)
        # 多种数据格式
        elif format in ['n', 'N', 'numpy']:
            return numpy.asarray(res)
        elif format in ['list', 'l', 'L']:
            return numpy.asarray(res).tolist()
        else:
            print(
                "DigQuant Error QA_fetch_stock_day format parameter %s is none of  \"P, p, pandas, pd , json, dict , n, N, numpy, list, l, L, !\" "
                % format
            )
            return None
    else:
        LOG.info(
            'DigQuant Error QA_fetch_stock_day data parameter start=%s end=%s is not right'
            % (start,
               end)
        )


def QA_fetch_stock_min_adv(
    code,
    start,
    end=None,
    frequence='1min',
    if_drop_index=True,
):
    """
    获取股票分钟线

    :param code:  字符串str eg 600085
    :param start: 字符串str 开始日期 eg 2011-01-01
    :param end:   字符串str 结束日期 eg 2011-05-01
    :param frequence: 字符串str 分钟线的类型 目前仅支持 1min 1m
    :param if_drop_index: Ture False ， dataframe drop index or not
    :return: QA_DataStruct_Stock_min 类型
    """
    if frequence in ['1min', '1m']:
        frequence = '1min'
    else:
        print(
            "DigQuant Error QA_fetch_stock_min_adv parameter frequence=%s is none of 1min 1m"
            % frequence
        )
        return None

    end = start if end is None else end

    if start == end:
        print(
            "DigQuant Error QA_fetch_stock_min_adv parameter code=%s , start=%s, end=%s is equal, should have time span! "
            % (code,
               start,
               end)
        )
        return None

    res = QA_fetch_stock_min(code, start, end, format='pd', frequence=frequence)
    if res is None:
        print(
            "DigQuant Error QA_fetch_stock_min_adv parameter code=%s , start=%s, end=%s frequence=%s call QA_fetch_stock_min return None"
            % (code,
               start,
               end,
               frequence)
        )
        return None
    else:
        res_set_index = res.set_index(['datetime', 'code'], drop=if_drop_index)
        # if res_set_index is None:
        #     print("QA Error QA_fetch_stock_min_adv set index 'datetime, code' return None")
        #     return None
        return QA_DataStruct_Stock_min(res_set_index)


def QA_fetch_stock_min(
        code,
        start,
        end,
        format='numpy',
        frequence='1min'
):
    # 获取股票分钟线
    if frequence in ['1min', '1m']:
        frequence = '1min'
    else:
        print(
            "DigQuant Error QA_fetch_stock_min parameter frequence=%s is none of 1min 1m"
            % frequence
        )

    _data = []
    # code checking
    code = QA_util_code_tolist(code)
    mkt_code = ['sse' if x[0] == '6' else 'szse' for x in code]
    code_df = pd.DataFrame({'code': code, 'market': mkt_code})
    code_dict = code_df.to_dict(orient='records')

    res = get_dqt_data_by_api(data_type='quote/stock/kdata/min',
                              input_data={
                                  "dateDurationTO": {
                                      "beginDate": start,
                                      "endDate": end
                                  },
                                  "filledUp": True,
                                  "marketCodeTOS": code_dict,
                                  "restorationWay": "NA"
                              })
    if res is None:
        pass
    else:
        res_list = []
        for tmp in res:
            res_list.extend(tmp['data'])
        res = pd.DataFrame(res_list)
        res = res[['code', 'tradeTime', 'open', 'high', 'low', 'close', 'volume', 'turnOver']]
        res = res.rename(columns={'tradeTime': 'datetime', 'turnOver': 'amount'})
        res['type'] = '1min'

    try:
        res = res.assign(
            datetime=pd.to_datetime(res.datetime).dt.tz_localize(None).dt.tz_localize('Asia/Shanghai')
        ).query('volume>1').drop_duplicates(['datetime',
                                             'code']).set_index(
                                                 'datetime',
                                                 drop=False
                                             )
        if res.empty:
            res = None
    except:
        res = None

    if format in ['P', 'p', 'pandas', 'pd']:
        return res
    elif format in ['json', 'dict']:
        return QA_util_to_json_from_pandas(res)
    # 多种数据格式
    elif format in ['n', 'N', 'numpy']:
        return numpy.asarray(res)
    elif format in ['list', 'l', 'L']:
        return numpy.asarray(res).tolist()
    else:
        print(
            "DigQuant Error QA_fetch_stock_min format parameter %s is none of  \"P, p, pandas, pd , json, dict , n, N, numpy, list, l, L, !\" "
            % format
        )
        return None


if __name__ == '__main__':
    aaa = get_dqt_data_by_api(data_type='quote/stock/kdata/day',
                              input_data={
                                  "dateDurationTO": {
                                      "beginDate": "2020-09-01",
                                      "endDate": "2020-09-05"
                                  },
                                  "filledUp": True,
                                  "marketCodeTOS": [
                                      {
                                          "code": "000001",
                                          "market": "szse"
                                      },
                                      {
                                          "code": "000002",
                                          "market": "szse"
                                      }
                                  ],
                                  "restorationWay": "NA"
                              })

    bbb = QA_fetch_stock_day(['000001', '000002'], '2020-09-01', '2020-09-05', 'pd')
    ccc = QA_fetch_stock_min(['000001', '000002'], '2020-11-11', '2020-11-12', 'pd')
